from flask import Flask, jsonify, request
import json
import random
import numpy
import torch
from ner.ner_case import NerCase
from utils.common_utils import get_parameters

ner = None
model_file_dict = None
app = Flask('医学命名实体识别服务')


@app.route('/predict', methods=['post', 'get'])
def predict():
    bl, params = get_parameters(request, 'contents')
    if not bl:
        return params
    contents = params[0]
    bl, params = get_parameters(request, 'mod_name')
    global ner
    if bl:
        mod_name = params[0]
        if not model_file_dict.__contains__(mod_name):
            return '模型不存在'
        ner = NerCase.load_from(model_file_dict[mod_name])
    contents = [x for x in contents if x.strip()]
    print('---- 收到一批数据 ----')
    print(contents[0])
    print('......', len(contents)-2)
    print(contents[-1])
    result = {
        'status': 'success',
        'message': 'success',
        'result': ner.predict_batch(contents)
    }
    return jsonify(result)


if __name__ == '__main__':
    print('------')

    # 固定随机数
    random.seed(1)
    numpy.random.seed(1)
    torch.manual_seed(1)
    torch.cuda.empty_cache()

    # 载入服务的设置
    print('载入服务配置')
    server_config = json.load(open('server_config.json', 'r', encoding='utf8'))
    model_file_dict = server_config['model_file_name']
    port = server_config['port']
    app.config['JSON_AS_ASCII'] = False

    # for key in model_file_dict.keys():
    key = 'roberta_wwm-dynamicfusion1x-strokes-wubi-zhengma-glyph-pinyin-bilstm1x-crf'
    # 载入模型
    print('初始化模型')
    ner = NerCase.load_from(model_file_dict[key]) # 从文件名载入已训练的模型
    ner.predict("肝脏肾脏 都是内脏。") # 让模型运行一次以热身
    # 启动服务
    print('启动服务')
    print(f'模型 {key}')
    #app.run(debug=True, port=port, use_reloader=False)
    app.run(host='0.0.0.0', port=port, use_reloader=False)


 # "model_dict": {
 #      RoBERT-WWM+BiLSTM+CRF: roberta_wwm-bilstm1x-crf
 #      RoBERT-WWM+BiLSTM+CRF+动态融合: roberta_wwm-dynamicfusion1x-bilstm1x-crf,
 #      RoBERT-WWM+BiLSTM+CRF+动态融合+笔画+五笔+郑码+汉字图像+拼音: roberta_wwm-dynamicfusion1x-strokes-wubi-zhengma-glyph-pinyin-bilstm1x-crf,
 #      RoBERT-WWM+BiLSTM+CRF+笔画: roberta_wwm-strokes-bilstm1x-crf,
 #      RoBERT-WWM+BiLSTM+CRF+拼音+笔画+五笔+郑码+汉字图像: roberta_wwm-strokes-wubi-zhengma-glyph-pinyin-bilstm1x-crf,
 #      RoBERT-WWM+BiLSTM+CRF+五笔: roberta_wwm-wubi-bilstm1x-crf
 #    }